using System;
using System.Linq;
namespace RayDen.Library.Core
{
    public class Distribution1D
    {

        internal float[] func, cdf;
        internal float funcInt;
        internal int count;


        public Distribution1D(float[] f)
        {
            var n = f.Length;
            this.count = n;
            this.func = new float[n];
            Array.Copy(f, this.func, n);
            this.cdf = new float[n+1];
            // Compute integral of step function at $x_i$
            this.cdf[0] = 0f;
            for(int i = 1; i < this.count+1; ++i)
                this.cdf[i] = this.cdf[i-1] + this.func[i-1] / n;

            // Transform step function integral into CDF
            this.funcInt = this.cdf[this.count];
            if(this.funcInt == 0f)
            {
                for(int i = 1; i < n+1; ++i)
                    this.cdf[i] = (float)i / (float)(n);
            }
            else
            {
                for(int i = 1; i < n+1; ++i)
                    this.cdf[i] /= this.funcInt;
            }
        }

        public float SampleContinuous(float u, out float pdf, out int off)
        {
            // Find surrounding CDF segments and _offset_
            var ptr = this.cdf.SkipWhile(item => item <= u).ToArray().Length;
            int offset = Math.Max(0, (int)(ptr-1));
            off = offset;
            //Assert(offset < count);
            //Assert(u >= cdf[offset] && u < cdf[offset+1]);

            // Compute offset along CDF segment
            float du = (u - this.cdf[offset]) / (this.cdf[offset+1] - this.cdf[offset]);
            //Assert(!isnan(du));

            // Compute PDF for sampled offset
            pdf = this.func[offset] / this.funcInt;

            // Return $x\in{}[0,1)$ corresponding to sample
            return (offset + du) / this.count;
        }

        public int SampleDiscrete(float u, out float pdf)
        {
            // Find surrounding CDF segments and _offset_
            var ptr = this.cdf.SkipWhile(item => item <= u).ToArray().Length;
            int offset = Math.Max(0, (int)(ptr-1));
            pdf = this.func[offset] / (this.funcInt * this.count);
            return offset;
        }
    }
}